import struct
import time
import socket


# 自定义异常类，表示客户端主动断开连接
class ClientDisconnectedError(Exception):
    """当客户端主动断开连接时抛出此异常"""
    pass


class FileTransferProtocol:
    def __init__(self):
        pass

    def pack_data(self, message_type: int, data: bytes) -> bytes:
        """
        打包数据为指定格式
        格式: | Type(2B) | Length(4B) | Data(NB) |

        :param message_type: 消息类型
        :param data: 要发送的数据
        :return: 打包后的字节数据
        """
        # 打包消息类型（2字节）
        type_header = struct.pack('H', message_type)
        # 打包数据长度（4字节）
        length_header = struct.pack('I', len(data))
        # 组合完整消息
        return type_header + length_header + data

    def receive_file(self, client_socket: socket.socket) -> tuple[int, bytes]:
        # 接收数据类型（2字节）+ 长度头（4字节）+ 获取实际数据长度
        data_type = self._recvall(2, client_socket)
        message_type = struct.unpack('H', data_type)[0]
        header = self._recvall(4, client_socket)
        length = struct.unpack('I', header)[0]
        data = self._recvall(length, client_socket)
        return message_type, data

    def _recvall(self, n, client_socket: socket.socket):
        """
        确保接收指定长度的字节数据
        :param client_socket: 客户端socket
        :param n: 需要接收的字节数
        :return: 接收到的字节数据
        :raises ClientDisconnectedError: 当客户端主动断开连接时
        :raises Exception: 其他异常情况
        """
        data = b''

        try:
            while len(data) < n:
                # 接收剩余所需字节数
                packet = client_socket.recv(n - len(data))
                if not packet:
                    raise ClientDisconnectedError("客户端主动断开连接")
                data += packet
            return data
        except Exception:
            # 其他异常向上抛出
            raise
